import requests
import json

def get_embedding(text):
    # Ollama 嵌入接口
    url = "http://127.0.0.1:11434/api/embeddings"
    payload = {
        "model": "nomic-embed-text:latest",  # 必须和 curl 返回的模型名一致
        "prompt": text  # 要嵌入的文本，支持单句或列表
    }
    headers = {"Content-Type": "application/json"}

    try:
        # 发送 POST 请求
        response = requests.post(url, headers=headers, data=json.dumps(payload))
        response.raise_for_status()  # 触发 HTTP 错误（如 404、500）
        result = response.json()
        return result["embedding"]  # 返回嵌入向量
    except requests.exceptions.HTTPError as e:
        print(f"HTTP 错误: {e}")
        print("响应内容:", response.text)  # 打印错误详情，方便排查
    except Exception as e:
        print(f"其他错误: {e}")
    return None

# 测试调用
if __name__ == "__main__":
    test_text = "这是一段测试嵌入的文本"
    embedding = get_embedding(test_text)
    if embedding:
        print(f"嵌入向量长度: {len(embedding)}")
        print("嵌入向量前 10 位:", embedding[:10])